import sys
import numpy as np
import matplotlib.pyplot as plt
import glob

# Make sure that caffe is on the python path:
caffe_root = '/home/cgf/caffeLabUsing/'  # this file is expected to be in {caffe_root}/examples
import sys
sys.path.insert(0, caffe_root + 'python')

import caffe

def reconginze(modelName, imgNames):
	import os
	rightClassFileName="rightClass.txt"
	if os.path.isfile(rightClassFileName):
		os.remove(rightClassFileName)
	
	rightClassFile = open("rightClass.txt", "wb")
	
	plt.rcParams['figure.figsize'] = (10, 10)
	plt.rcParams['image.interpolation'] = 'nearest'
	plt.rcParams['image.cmap'] = 'gray'

	if not os.path.isfile(caffe_root + 'models/' + modelName + '/caffe_train.caffemodel'):
		print("Downloading pre-trained CaffeNet model...")
		sys.exit()
	
	caffe.set_mode_cpu()
	net = caffe.Net(caffe_root + 'models/' + modelName + '/deploy.prototxt', caffe_root + 'models/' + modelName + '/caffe_train.caffemodel', caffe.TEST)

	# input preprocessing: 'data' is the name of the input blob == net.inputs[0]
	transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
	transformer.set_transpose('data', (2,0,1))
	transformer.set_mean('data', np.load(caffe_root + 'python/caffe/' + modelName + '/' + modelName + '_mean.npy').mean(1).mean(1)) # mean pixel
	transformer.set_raw_scale('data', 255)  # the reference model operates on images in [0,255] range instead of [0,1]
	transformer.set_channel_swap('data', (2,1,0))  # the reference model has channels in BGR order instead of RGB

	# load labels
	model_labels_filename = caffe_root + 'data/' + modelName +'/synset_words.txt'
	labels = np.loadtxt(model_labels_filename, str, delimiter='\t')

	# set net to batch size of 50
	net.blobs['data'].reshape(50,3,227,227)
	filePath=[]
	dirNum = 0
	totalValNum = 0
	totalRightNum = 0
    	if(imgNames.endswith(".jpg") == True):
        	print "start file scan"
		filePath.append(imgNames)
    	else:
        	print "start dir scan"
		for dir in os.listdir(imgNames):
			dirNum += 1
            		for file in glob.glob(imgNames + "/" + dir + "/" + dir + "_*.jpg"):
				#print str(dirNum) + "," + dir + "," + file
                		filePath.append(file)
				totalValNum += 1
	print "totalValNum:" + str(totalValNum)
	top_k = 0
	for imgName in filePath:
		tag = imgName.rsplit('/')[-2]
		#print "imgName=" + imgName
		net.blobs['data'].data[...] = transformer.preprocess('data', caffe.io.load_image(imgName))
		out = net.forward()
		result="Predicted class is #{}.".format(out['prob'][0].argmax())
		#print(result)

    		#for showing image
		#plt.imshow(transformer.deprocess('data', net.blobs['data'].data[0]))
		#plt.show()


		# sort top k predictions from softmax output
		top_k = net.blobs['prob'].data[0].flatten().argsort()[-1:-6:-1]
		#print labels[top_k]
		labelsLen = len(labels[top_k])
		#print labelsLen
		isReconginze = False
		for index in range(labelsLen):
			labelName = labels[top_k][index]
			if labels[top_k][0].find(tag) >= 0:
				isReconginze = True
		if isReconginze == True:
			rightClassFile.write(imgName + "\n")
			totalRightNum += 1
			print "the same as the tag, totalRightNum:" + str(totalRightNum)
		else:
			print "the diff as the tag, totalRightNum:" + str(totalRightNum)
	if totalValNum > 0:
		rightClassFile.write("totalValNum:" + str(totalValNum) + ",totalRightNum:" + str(totalRightNum) + ",the right rate:" + str((totalRightNum * 1.0) / (totalValNum * 1.0)) + "\n")
		print "totalValNum:" + str(totalValNum) + ",totalRightNum:" + str(totalRightNum) + ",the right rate:" + str((totalRightNum * 1.0) / (totalValNum * 1.0))
	
	#print labels[top_k]
	#else:
		#print "wrong path"
		#rightClassFile.close()
		#return "wrong path"
	rightClassFile.close()
	return labels[top_k]

if __name__=='__main__':
	if len(sys.argv) != 3:
		print "Usage: python caffeRecongination.py modelName{such as oxfordCatsAndDogs} imgNames{cat.jpg}"
		sys.exit()

	modelName=sys.argv[1]
	imgNames=sys.argv[2]
	reconginze(modelName, imgNames)
	#imgName="/home/cgf/Abyssinian_1.jpg"
